import numpy as np
import joblib
from scipy.fft import fft, fftfreq


class RunPython:
    def __init__(self, model_path, fs):
        self.bad_channels = [1, 4, 5, 6, 8, 11, 12, 15, 17]
        self.fs = fs
        self.model = joblib.load(model_path)

    def extract_features(self, ecog_data, f_max=200, f_split=140, low_step=20, high_step=30):
        """
        胡炫烨模型特征提取方法
        :param ecog_data:
        :param f_max:
        :param f_split:
        :param low_step:
        :param high_step:
        :return:
        """
        # 剔除坏道
        data = np.delete(ecog_data, self.bad_channels, axis=1)
        # 参数预设
        n_samples, n_channels = data.shape

        # 构造所有频段 [(f_start, f_end), ...]
        bands = [(i, i + low_step) for i in range(0, f_split, low_step)] + \
                [(i, i + high_step) for i in range(f_split, f_max, high_step)]
        n_bands = len(bands)

        feature_list = []
        for ch in range(n_channels):
            sig = data[:, ch]

            # 计算 FFT
            fft_vals = np.fft.fft(sig)
            freqs = np.fft.fftfreq(len(sig), 1 / self.fs)  # 计算频率轴

            # 只取正频率部分
            positive_freqs = freqs[freqs >= 0]
            positive_fft_vals = np.abs(fft_vals[freqs >= 0])  # FFT 的幅值

            # 提取各频段平均功率
            for f_start, f_end in bands:
                # 找出频率索引范围
                idx = np.where((positive_freqs >= f_start) & (positive_freqs < f_end))[0]
                fft_mean = np.mean(positive_fft_vals[idx]) if len(idx) > 0 else 0
                feature_list.append(fft_mean)

            # 提取RMS特征
            rms_value = np.sqrt(np.mean(sig ** 2))
            feature_list.append(rms_value)
        # 标准化
        # scaler = StandardScaler()
        # X_features = scaler.fit_transform(np.array(feature_list).reshape(1, -1))
        return np.array(feature_list).reshape(1, -1)

    def run_model(self, data):
        """
        返回1的score
        :param data:
        :return:
        """
        X_features = self.extract_features(data)
        y_prob = self.model.predict_proba(X_features)[:, 1]  # 获取1的概率
        # print("1概率：", y_prob[0])
        return y_prob[0]

    def close(self):
        print("close")
